"""
This file provides functions for calculating pvalues using dynamic programming
"""

from math import *
import array
import numpy as np
import random


def compute_dynamic_pval(p, d, smooth_pvalue = False, \
                         steps_per_position = 1000): 
    """
    Input: 
    p - list of the probabilities of all the peptides of this protein
    d - list including 1 for peptides matched to MS1-features, 0 otherwise
    smooth_pvalue - True if the pvalues should be smoothed 
    steps_per_position - discretization step

    Output: estimated p value
    """

    # calculate the vector x and the sum of log(pi)
    x = []
    sum_logp = 0.0
    for p_i in p:
        x_i = log( (1.0 - p_i) / p_i )
        x.append( x_i )
        sum_logp += log(p_i)

    # discretize each xi to get li and compute the sum of all li 
    l = []  
    sum_l = np.longdouble(0)
    k = max(x) / steps_per_position
    for x_i in x:
        l_i = int( round(x_i / k) )
        l.append(l_i)
        sum_l += l_i

    # express P(D|R = 0) in terms of li (D = obs config)
    sum_thresh = np.longdouble(0)
    n=len(p)
    for i in range(n):
        if d[i] == 0:
            sum_thresh += l[i]

    # dynamic programming, the easy way  
    f = array.array('d', (0.0,) * (sum_thresh + 1))
    f[0] = 1
    c=0.0
    for j in xrange(len(l)):
        if j%30==29:
            f_m=max(f)
            c+=log(f_m)
            for i in xrange(len(f)):
                f[i]/=f_m
        for i in xrange(sum_thresh, l[j] - 1, -1):
            f[i] += f[i-l[j]]

    # calculate the final p-value 
    sum_prob = 0.0
    for i in xrange(sum_thresh):
        if f[i]==0:
            continue
        sum_prob += exp( i * k + log(f[i]) + sum_logp + c)

    # add the last term separately
    if smooth_pvalue:
       rd = random.random()        
    else:
       rd = 1.0
    i = int(round(sum_thresh))
    sum_prob += exp( i * k + log(f[i]) + sum_logp + c) * rd 
 
    # rounding errors may lead to slightly larger values than 1.0        
    if sum_prob > 1.0:
       sum_prob = 1.0 

    return sum_prob


